import cv2
import numpy as np
import textwrap

# 确保 calib_params.npz 文件存在且包含 mtx 和 dist
# 为了让脚本可以独立运行，我们在这里创建一个虚拟的参数文件
def create_dummy_params():
    # 典型广角镜头的参数
    mtx = np.array([
        [800.0, 0.0,   512.0],
        [0.0,   800.0, 384.0],
        [0.0,   0.0,   1.0]
    ], dtype=np.float64)
    # 典型畸变参数
    dist = np.array([-0.4, 0.2, 0.0, 0.0, 0.0], dtype=np.float64)
    np.savez('calib_params.npz', mtx=mtx, dist=dist)
    print("Created dummy 'calib_params.npz' for demonstration.")

try:
    data = np.load('calibration/calib_params.npz')
    mtx = data['mtx']
    dist = data['dist']
except FileNotFoundError:
    create_dummy_params()
    data = np.load('calib_params.npz')
    mtx = data['mtx']
    dist = data['dist']


# --- 1. 用户配置区 ---
# 请在此处填入你的实际参数

# 原始高分辨率 (标定时使用的分辨率)
ORIG_WIDTH = 1024
ORIG_HEIGHT = 768

# 目标低分辨率 (最终处理时使用的分辨率)
TARGET_WIDTH = 640
TARGET_HEIGHT = 480

# 原始高分辨率下的相机内参矩阵 (mtx)
MTX_ORIG = mtx

# 畸变系数 (dist) - 这个与分辨率无关
DIST_COEFFS = dist


# --- 2. 计算与代码生成 ---

def calculate_and_generate_code():
    """
    计算缩放后的参数并生成使用 remap 的优化C++代码
    """
    print("--- [Step 1] Scaling camera matrix...")
    # 复制原始矩阵以进行修改
    mtx_scaled = MTX_ORIG.copy()

    # 计算缩放比例
    scale_x = TARGET_WIDTH / float(ORIG_WIDTH)
    scale_y = TARGET_HEIGHT / float(ORIG_HEIGHT)

    # 缩放内参矩阵的 fx, fy, cx, cy
    mtx_scaled[0, 0] *= scale_x  # fx
    mtx_scaled[1, 1] *= scale_y  # fy
    mtx_scaled[0, 2] *= scale_x  # cx
    mtx_scaled[1, 2] *= scale_y  # cy
    
    print(f"Original Mtx:\n{MTX_ORIG}")
    print(f"Scaled Mtx for {TARGET_WIDTH}x{TARGET_HEIGHT}:\n{mtx_scaled}\n")

    print("--- [Step 2] Calculating optimal new matrix and ROI for target resolution...")
    # 使用缩放后的参数计算新的最佳相机矩阵和ROI
    # alpha=1 表示保留所有原始像素，可能会有黑边
    new_mtx_scaled, roi = cv2.getOptimalNewCameraMatrix(
        cameraMatrix=mtx_scaled,
        distCoeffs=DIST_COEFFS,
        imageSize=(TARGET_WIDTH, TARGET_HEIGHT),
        alpha=1,
        newImgSize=(TARGET_WIDTH, TARGET_HEIGHT)
    )
    
    # ROI (x, y, w, h)
    roi_x, roi_y, roi_w, roi_h = roi

    print(f"Optimal New Mtx for {TARGET_WIDTH}x{TARGET_HEIGHT}:\n{new_mtx_scaled}")
    print(f"ROI for cropping: x={roi_x}, y={roi_y}, w={roi_w}, h={roi_h}\n")

    print("--- [Step 3] Generating Optimized C++ Code (using initUndistortRectifyMap and remap) ---\n")
    generate_cpp_remap(mtx_scaled, DIST_COEFFS, new_mtx_scaled, roi)


def format_mat_for_cpp(mat, name, dtype="double"):
    """将numpy矩阵格式化为C++ cv::Mat初始化代码"""
    rows, cols = mat.shape
    # 确保浮点数格式正确
    mat_data = ", ".join(f"{x:.10f}" for x in mat.flatten())
    return f"const cv::Mat {name} = (cv::Mat_<{dtype}>({rows}, {cols}) << {mat_data});"


def generate_cpp_remap(mtx, dist, new_mtx, roi):
    """
    生成使用 initUndistortRectifyMap 和 remap 的优化C++代码。
    """
    
    cpp_code = f"""
// =================================================================================
//   (Optimized) Pre-computed Camera Calibration Code for {TARGET_WIDTH}x{TARGET_HEIGHT} Resolution
//   Generated by Python script. Uses initUndistortRectifyMap and remap for high performance.
// =================================================================================

#include <opencv2/opencv.hpp>

// --- Generated Constants ---

// Scaled camera matrix for {TARGET_WIDTH}x{TARGET_HEIGHT} input images.
{format_mat_for_cpp(mtx, "PRECOMPUTED_MTX_LOW_RES")}

// Distortion coefficients (resolution-independent).
{format_mat_for_cpp(dist, "PRECOMPUTED_DIST")}

// Optimal new camera matrix for undistortion at {TARGET_WIDTH}x{TARGET_HEIGHT}.
{format_mat_for_cpp(new_mtx, "PRECOMPUTED_NEW_MTX_LOW_RES")}

// Region of Interest (ROI) for cropping black borders at {TARGET_WIDTH}x{TARGET_HEIGHT}.
const cv::Rect PRECOMPUTED_ROI_LOW_RES({roi[0]}, {roi[1]}, {roi[2]}, {roi[3]});


/**
 * @brief A class to handle high-performance image undistortion.
 *
 * It pre-computes the rectification maps in its constructor ONCE.
 * The `undistort()` method then uses these maps to process images very quickly.
 * This is ideal for video processing.
 */
class PrecomputedUndistorter
{{
public:
    /**
     * @brief Constructor that prepares the undistortion maps.
     *
     * This is the only "slow" part and should be called only once.
     */
    PrecomputedUndistorter()
        : m_roi(PRECOMPUTED_ROI_LOW_RES)
    {{
        const cv::Size image_size({TARGET_WIDTH}, {TARGET_HEIGHT});

        // Pre-calculate the rectification maps.
        // CV_16SC2 is a more compact and faster map format than the default CV_32FC1.
        cv::initUndistortRectifyMap(
            PRECOMPUTED_MTX_LOW_RES,
            PRECOMPUTED_DIST,
            cv::Mat(), // Optional rectification (not needed for monocular)
            PRECOMPUTED_NEW_MTX_LOW_RES,
            image_size,
            CV_16SC2, // Map type for remap
            m_map1,
            m_map2
        );
    }}

    /**
     * @brief (Highly Optimized) Undistorts and crops an image using pre-computed maps.
     *
     * @param img_low_res The distorted input image, which MUST be {TARGET_WIDTH}x{TARGET_HEIGHT}.
     * @return cv::Mat The undistorted and cropped image.
     */
    cv::Mat undistort(const cv::Mat& img_low_res) const
    {{
        // 1. Apply the rectification map. This is much faster than cv::undistort().
        cv::Mat dst;
        cv::remap(img_low_res, dst, m_map1, m_map2, cv::INTER_LINEAR);

        // 2. Crop to the valid region using the pre-calculated ROI.
        // .clone() creates a deep copy, making the returned Mat independent.
        cv::Mat cropped_dst = dst(m_roi).clone();

        return cropped_dst;
    }}

private:
    cv::Mat m_map1, m_map2;
    cv::Rect m_roi;
}};


// ======================= EXAMPLE USAGE (main function) ==========================
/*
#include <iostream>

// (Paste the generated code block, including the class, above this main function)

int main()
{{
    // 1. Create an instance of the undistorter.
    // The constructor will be called here, computing the maps once.
    PrecomputedUndistorter undistorter;
    std::cout << "Undistortion maps have been pre-computed." << std::endl;
    
    // 2. Read a high-resolution distorted image.
    // Replace with your actual image path.
    cv::Mat high_res_img = cv::imread("your_distorted_image_1024x768.jpg");
    if (high_res_img.empty())
    {{
        std::cerr << "Error: Could not read the image." << std::endl;
        return -1;
    }}

    // 3. Resize the image to the target resolution FIRST.
    // This is the key to performance.
    cv::Mat low_res_img;
    cv::resize(high_res_img, low_res_img, cv::Size({TARGET_WIDTH}, {TARGET_HEIGHT}), 0, 0, cv::INTER_LINEAR);

    // 4. Call the highly optimized undistortion method.
    // This can be done in a loop for a video stream and will be very fast.
    cv::Mat final_result = undistorter.undistort(low_res_img);

    // 5. Display the results.
    cv::imshow("Original High-Res Image", high_res_img);
    cv::imshow("Resized (but still distorted) Low-Res Image", low_res_img);
    cv::imshow("Final Undistorted and Cropped Result", final_result);

    std::cout << "Process finished. Final image size: " << final_result.size() << std::endl;
    cv::waitKey(0);

    return 0;
}}
*/
"""
    # 使用textwrap.dedent来移除由于f-string引起的多余缩进
    print(textwrap.dedent(cpp_code).strip())


# --- 3. 运行脚本 ---
if __name__ == "__main__":
    calculate_and_generate_code()


# 保存标定参数为XML文件
# fs = cv2.FileStorage('calibration_parameters.xml', cv2.FILE_STORAGE_WRITE)
# fs.write('cameraMatrix', mtx)
# fs.write('distCoeffs', dist)
# fs.release()
# print("Camera calibration parameters saved to calibration_parameters.xml")



# # 读一张要校正的照片
# img = cv2.imread('./calib_imgs/calib_01.jpg')  # 换成你实际图片路径
# h, w = img.shape[:2]

# # 计算新的校正映射
# newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w,h), 1, (w,h))

# # 去畸变
# dst = cv2.undistort(img, mtx, dist, None, newcameramtx)

# # 裁剪有效区域（可选）
# x, y, w, h = roi
# dst = dst[y:y+h, x:x+w]

# cv2.imshow('Undistorted', dst)
# cv2.waitKey(0)
# cv2.destroyAllWindows()
# cv2.imwrite('./undistorted_01.jpg', dst)
